package com.robot.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.SendTo;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Controller;

import com.robot.service.LangChainService;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * WebSocket控制器，处理实时消息通信
 */
@Controller
public class WebSocketController {

    private static final Logger logger = LoggerFactory.getLogger(WebSocketController.class);

    @Autowired
    private LangChainService langChainService;

    @Autowired
    private SimpMessagingTemplate messagingTemplate;
    
    @Autowired
    private ObjectMapper objectMapper;

    /**
     * 处理自然语言查询，并将处理过程实时推送给客户端
     */
    @MessageMapping("/query")
    @SendTo("/topic/response")
    public Map<String, Object> processQuery(Map<String, String> request) {
        Map<String, Object> response = new HashMap<>();
        
        String sessionId = request.getOrDefault("sessionId", "unknown-session");
        String query = request.get("query");
        
        if (query == null || query.trim().isEmpty()) {
            response.put("success", false);
            response.put("message", "查询不能为空");
            return response;
        }
        
        // 异步处理查询
        CompletableFuture.runAsync(() -> {
            try {
                // 发送开始处理的状态更新
                sendStatusUpdate(sessionId, "开始处理查询: " + query);
                
                // 处理查询
                String result = langChainService.processQuery(query);
                
                // 构建成功响应
                Map<String, Object> successResponse = new HashMap<>();
                successResponse.put("success", true);
                successResponse.put("response", result);
                successResponse.put("sessionId", sessionId);
                
                // 发送最终结果
                messagingTemplate.convertAndSend("/topic/response/" + sessionId, successResponse);
                
            } catch (Exception e) {
                logger.error("处理查询失败", e);
                
                // 构建错误响应
                Map<String, Object> errorResponse = new HashMap<>();
                errorResponse.put("success", false);
                errorResponse.put("message", "处理查询失败: " + e.getMessage());
                errorResponse.put("sessionId", sessionId);
                
                // 发送错误结果
                messagingTemplate.convertAndSend("/topic/response/" + sessionId, errorResponse);
            }
        });
        
        // 立即返回确认消息
        response.put("success", true);
        response.put("message", "查询已接收，正在处理");
        response.put("sessionId", sessionId);
        
        return response;
    }
    
    /**
     * 处理聊天消息，并将处理过程实时推送给客户端
     */
    @MessageMapping("/chat")
    @SendTo("/topic/chat")
    public Map<String, Object> processChat(Map<String, String> request) {
        Map<String, Object> response = new HashMap<>();
        
        String sessionId = request.getOrDefault("sessionId", "unknown-session");
        String message = request.get("message");
        
        if (message == null || message.trim().isEmpty()) {
            response.put("success", false);
            response.put("message", "消息不能为空");
            return response;
        }
        
        // 异步处理聊天
        CompletableFuture.runAsync(() -> {
            try {
                // 发送开始处理的状态更新
                sendStatusUpdate(sessionId, "正在处理聊天消息");
                
                // 处理聊天
                String result = langChainService.chat(message);
                
                // 构建成功响应
                Map<String, Object> successResponse = new HashMap<>();
                successResponse.put("success", true);
                successResponse.put("response", result);
                successResponse.put("sessionId", sessionId);
                
                // 发送最终结果
                messagingTemplate.convertAndSend("/topic/chat/" + sessionId, successResponse);
                
            } catch (Exception e) {
                logger.error("处理聊天失败", e);
                
                // 构建错误响应
                Map<String, Object> errorResponse = new HashMap<>();
                errorResponse.put("success", false);
                errorResponse.put("message", "处理聊天失败: " + e.getMessage());
                errorResponse.put("sessionId", sessionId);
                
                // 发送错误结果
                messagingTemplate.convertAndSend("/topic/chat/" + sessionId, errorResponse);
            }
        });
        
        // 立即返回确认消息
        response.put("success", true);
        response.put("message", "消息已接收，正在处理");
        response.put("sessionId", sessionId);
        
        return response;
    }
    
    /**
     * 发送状态更新消息
     */
    private void sendStatusUpdate(String sessionId, String status) {
        Map<String, Object> statusUpdate = new HashMap<>();
        statusUpdate.put("type", "status");
        statusUpdate.put("status", status);
        statusUpdate.put("sessionId", sessionId);
        
        messagingTemplate.convertAndSend("/topic/status/" + sessionId, statusUpdate);
    }
} 